#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2
import sensor_msgs_py.point_cloud2 as pc2
import numpy as np
import threading
from std_srvs.srv import Trigger
import os
from datetime import datetime
import time

class PointCloudSaver(Node):
    def __init__(self):
        # 先等待一下确保ROS 2完全初始化
        time.sleep(2)
        
        super().__init__('pointcloud_saver')
        
        self.get_logger().info('点云保存节点初始化...')
        
        # 订阅RTAB-Map的点云话题
        self.subscription = self.create_subscription(
            PointCloud2,
            '/rtabmap/cloud_map',
            self.pointcloud_callback,
            10)
        
        self.pointcloud_data = None
        self.lock = threading.Lock()
        
        self.get_logger().info('点云保存节点已启动 (修正版)')
        
        # 创建保存服务
        self.save_service = self.create_service(
            Trigger,
            'save_pointcloud',
            self.save_callback)
        
        # 创建保存目录
        self.save_dir = os.path.expanduser('~/pointcloud_maps')
        os.makedirs(self.save_dir, exist_ok=True)
        self.get_logger().info(f'点云将保存到: {self.save_dir}')
    
    def pointcloud_callback(self, msg):
        with self.lock:
            try:
                # 转换点云数据
                points = pc2.read_points(msg, field_names=("x", "y", "z", "rgb"), skip_nans=True)
                points_list = list(points)
                
                if len(points_list) > 0:
                    # 转换为numpy数组
                    self.pointcloud_data = np.array(points_list)
                    self.get_logger().info(f'收到点云数据: {len(self.pointcloud_data)} 个点')
                else:
                    self.get_logger().warn('收到的点云数据为空')
                    
            except Exception as e:
                self.get_logger().error(f'处理点云数据时出错: {str(e)}')
    
    def save_callback(self, request, response):
        with self.lock:
            if self.pointcloud_data is not None and len(self.pointcloud_data) > 0:
                try:
                    # 生成文件名
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    output_path = os.path.join(self.save_dir, f'office_map_{timestamp}.ply')
                    
                    # 手动创建PLY文件
                    self.save_ply_manual(output_path)
                    
                    response.success = True
                    response.message = f'点云已保存到: {output_path}，包含 {len(self.pointcloud_data)} 个点'
                    self.get_logger().info(response.message)
                    
                except Exception as e:
                    response.success = False
                    response.message = f'保存失败: {str(e)}'
                    self.get_logger().error(response.message)
            else:
                response.success = False
                response.message = '没有可用的点云数据'
                self.get_logger().warn(response.message)
        
        return response
    
    def save_ply_manual(self, filepath):
        """手动创建PLY文件，不依赖Open3D"""
        with open(filepath, 'w') as f:
            # 写入PLY文件头
            f.write("ply\n")
            f.write("format ascii 1.0\n")
            f.write(f"element vertex {len(self.pointcloud_data)}\n")
            f.write("property float x\n")
            f.write("property float y\n")
            f.write("property float z\n")
            f.write("property uchar red\n")
            f.write("property uchar green\n")
            f.write("property uchar blue\n")
            f.write("end_header\n")
            
            # 写入点云数据
            for point in self.pointcloud_data:
                x, y, z = point[0], point[1], point[2]
                
                # 处理颜色（如果有）
                if len(point) >= 4:
                    rgb = int(point[3])
                    r = (rgb >> 16) & 0xff
                    g = (rgb >> 8) & 0xff
                    b = rgb & 0xff
                else:
                    r, g, b = 255, 255, 255  # 默认白色
                
                # 写入点数据
                f.write(f"{x:.6f} {y:.6f} {z:.6f} {r} {g} {b}\n")

def main():
    try:
        rclpy.init()
        node = PointCloudSaver()
        
        try:
            node.get_logger().info('点云保存节点开始运行...')
            rclpy.spin(node)
        except KeyboardInterrupt:
            node.get_logger().info('节点被用户中断')
        finally:
            node.destroy_node()
            rclpy.shutdown()
    except Exception as e:
        print(f"节点启动失败: {e}")
        print("请确保:")
        print("1. ROS 2环境已设置: source /opt/ros/humble/setup.bash")
        print("2. 没有其他ROS 2节点冲突")

if __name__ == '__main__':
    main()
